from .data_manager_Market1501 import Market1501
from .data_manager_DukeMTMC import DukeMTMC
from .data_manager_DebugDataset import DebugDataset

__factory = {
    # debug test
    "DebugDataset": DebugDataset,

    # Market-1501-v15.09.15
    "Market-1501-v15.09.15": Market1501,
    "market1501": Market1501,
    "Market1501": Market1501,
    "market-1501": Market1501,
    "Market-1501": Market1501,
    "market": Market1501,
    "Market": Market1501,

    # DukeMTMC-reID
    "DukeMTMC-reID": DukeMTMC,
    "DukeMTMC": DukeMTMC,
    "duke": DukeMTMC,
    "Duke": DukeMTMC,

    # cuhk03
}


def dataset_names():
    return sorted(__factory.keys())


def get_dataset(dataset_name, dataset_container_dirpath, **kwargs):
    if dataset_name not in __factory:
        raise KeyError("Unknown Dataset:", dataset_name)
    return __factory[dataset_name](dataset_container_dirpath, **kwargs)


# "dataset_container_dirpath" is the path of the folder that contains all the datasets for image-based ReID altogether. The folder is recommended to be named as DATA4REID by the habit of the author ZHOU YUXUAN. In the folder, each famous dataset sub-folders should be arranged parallelly.
# ==============================================
# DATA4REID               (-- dataset_container)
# |-cuhk03                      (--dataset_name)
# |-DukeMTMC-reID               (--dataset_name)
# |-Market-1501-v15.09.15       (--dataset_name)
# ==============================================
